#!/usr/bin/env python3
"""
Tunnel Controller with Link Statistics Monitoring
"""

import argparse
import grpc
import os
import sys
from datetime import datetime
from pathlib import Path
from time import sleep
from typing import Dict, List, Tuple

# Import P4Runtime libraries
sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '../../utils/'))
import p4runtime_lib.bmv2
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections
import p4runtime_lib.helper


class TunnelMapper:
    """Manages tunnel mappings and routing configurations"""
    
    def __init__(self):
        self.host_port = 1
        # tunnel_id: (source_switch, dest_switch, output_port, dest_mac)
        self.tunnel_routes = {
            100: ('s1', 's2', 2, "08:00:00:00:02:22"),
            101: ('s2', 's1', 2, "08:00:00:00:01:11"),
            200: ('s1', 's3', 3, "08:00:00:00:03:33"),
            201: ('s3', 's1', 2, "08:00:00:00:01:11"),
            300: ('s2', 's3', 3, "08:00:00:00:03:33"),
            301: ('s3', 's2', 3, "08:00:00:00:02:22")
        }
        
        # (switch_name, dest_ip): tunnel_id
        self.ip_mappings = {
            ('s1', "10.0.2.2"): 100,
            ('s1', "10.0.3.3"): 200,
            ('s2', "10.0.1.1"): 101,
            ('s2', "10.0.3.3"): 300,
            ('s3', "10.0.1.1"): 201,
            ('s3', "10.0.2.2"): 301
        }
        
        # switch_name: (target_ip, reply_mac)
        self.arp_responses = {
            's1': ("10.0.1.10", "08:00:00:00:01:00"),
            's2': ("10.0.2.20", "08:00:00:00:02:00"),
            's3': ("10.0.3.30", "08:00:00:00:03:00")
        }


class StatsLogger:
    """Handles link statistics logging"""
    
    def __init__(self, log_directory: str = "logs"):
        self.log_dir = Path(log_directory)
        self.log_dir.mkdir(exist_ok=True)
        self.log_handles = {}
        self._init_logs()
    
    def _init_logs(self):
        """Initialize log files for each network link"""
        links = ['s1s2', 's1s3', 's2s3']
        for link in links:
            log_path = self.log_dir / f"{link}.txt"
            self.log_handles[link] = open(log_path, 'w', encoding='utf-8')
            self.log_handles[link].write(f"# {link} stats started at {datetime.now()}\n")
            self.log_handles[link].flush()
    
    def record_stats(self, counter_info: Dict):
        """Log current statistics for all links"""
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        tunnel_mapper = TunnelMapper()
        
        for link, tunnels in self._link_tunnels().items():
            entries = []
            
            for tunnel_id, direction in tunnels:
                if tunnel_id in counter_info:
                    src_sw, dst_sw = tunnel_mapper.tunnel_routes[tunnel_id][:2]
                    count = counter_info[tunnel_id]
                    
                    entries.append(f"{dst_sw} received {count} pkts (tunnel {tunnel_id})")
                    entries.append(f"{src_sw} sent {count} pkts (tunnel {tunnel_id})")
                    entries.append(f"Path: {src_sw}->{dst_sw}")
            
            if entries:
                log_file = self.log_handles[link]
                log_file.write(f"\n# {timestamp}\n")
                for idx, entry in enumerate(entries, 1):
                    log_file.write(f"{idx}. {entry}\n")
                log_file.flush()
    
    def _link_tunnels(self) -> Dict:
        """Mapping of links to their associated tunnels"""
        return {
            's1s2': [(100, 'in'), (101, 'in')],
            's1s3': [(200, 'in'), (201, 'in')],
            's2s3': [(300, 'in'), (301, 'in')]
        }
    
    def close_logs(self):
        """Clean up log file handles"""
        for handle in self.log_handles.values():
            handle.close()


class P4Controller:
    """Main controller class for managing P4 switches"""
    
    def __init__(self, p4info_helper, bmv2_json_path):
        self.p4_helper = p4info_helper
        self.bmv2_path = bmv2_json_path
        self.tunnel_mapper = TunnelMapper()
        self.stats_logger = StatsLogger()
        self.switch_connections = {}
    
    def connect_switches(self):
        """Establish connections to all switches"""
        switch_specs = [
            ('s1', '127.0.0.1:50051', 0),
            ('s2', '127.0.0.1:50052', 1),
            ('s3', '127.0.0.1:50053', 2)
        ]
        
        for name, addr, dev_id in switch_specs:
            self.switch_connections[name] = p4runtime_lib.bmv2.Bmv2SwitchConnection(
                name=name,
                address=addr,
                device_id=dev_id,
                proto_dump_file=f'logs/{name}-p4runtime-requests.txt'
            )
            self.switch_connections[name].MasterArbitrationUpdate()
            self.switch_connections[name].SetForwardingPipelineConfig(
                p4info=self.p4_helper.p4info,
                bmv2_json_file_path=self.bmv2_path
            )
            print(f"Initialized switch: {name}")
    
    def setup_forwarding(self):
        """Configure all forwarding rules on switches"""
        self._setup_ip_routes()
        self._setup_tunnel_rules()
        self._setup_arp_rules()
        print("All forwarding rules configured")
    
    def _setup_ip_routes(self):
        """Configure IPv4 routing rules"""
        for (sw_name, dst_ip), tunnel_id in self.tunnel_mapper.ip_mappings.items():
            table_entry = self.p4_helper.buildTableEntry(
                table_name="MyIngress.ipv4_lpm",
                match_fields={"hdr.ipv4.dstAddr": (dst_ip, 32)},
                action_name="MyIngress.myTunnel_ingress",
                action_params={"dst_id": tunnel_id}
            )
            self.switch_connections[sw_name].WriteTableEntry(table_entry)
    
    def _setup_tunnel_rules(self):
        """Configure tunnel forwarding rules"""
        for tid, (ingress_sw, egress_sw, port, dst_mac) in self.tunnel_mapper.tunnel_routes.items():
            # Forwarding entry
            forward_entry = self.p4_helper.buildTableEntry(
                table_name="MyIngress.myTunnel_exact",
                match_fields={"hdr.myTunnel.dst_id": tid},
                action_name="MyIngress.myTunnel_forward",
                action_params={"port": port}
            )
            self.switch_connections[ingress_sw].WriteTableEntry(forward_entry)
            
            # Egress entry
            egress_entry = self.p4_helper.buildTableEntry(
                table_name="MyIngress.myTunnel_exact",
                match_fields={"hdr.myTunnel.dst_id": tid},
                action_name="MyIngress.myTunnel_egress",
                action_params={
                    "dstAddr": dst_mac,
                    "port": self.tunnel_mapper.host_port
                }
            )
            self.switch_connections[egress_sw].WriteTableEntry(egress_entry)
    
    def _setup_arp_rules(self):
        """Configure ARP response rules"""
        for sw_name, (target_ip, reply_mac) in self.tunnel_mapper.arp_responses.items():
            table_entry = self.p4_helper.buildTableEntry(
                table_name="MyIngress.arp_match",
                match_fields={
                    "hdr.arp.oper": [1],
                    "hdr.arp.tpa": [target_ip, 32]
                },
                action_name="MyIngress.send_arp_reply",
                action_params={"macAddr": reply_mac}
            )
            self.switch_connections[sw_name].WriteTableEntry(table_entry)
    
    def get_counter_data(self) -> Dict[int, int]:
        """Collect counter statistics from all switches"""
        counters = {}
        
        for tid in self.tunnel_mapper.tunnel_routes.keys():
            ingress_sw, egress_sw = self.tunnel_mapper.tunnel_routes[tid][:2]
            
            ing_count = self._read_counter(
                self.switch_connections[ingress_sw], 
                "MyIngress.ingressTunnelCounter", tid
            )
            egr_count = self._read_counter(
                self.switch_connections[egress_sw],
                "MyIngress.egressTunnelCounter", tid
            )
            
            counters[tid] = ing_count if ing_count > 0 else egr_count
        
        return counters
    
    def _read_counter(self, switch, counter_name: str, index: int) -> int:
        """Read a specific counter value"""
        try:
            for response in switch.ReadCounters(
                self.p4_helper.get_counters_id(counter_name), index):
                for entity in response.entities:
                    return entity.counter_entry.data.packet_count
        except Exception:
            pass
        return 0
    
    def show_stats(self):
        """Display current statistics to console"""
        print(f"\nLink Statistics - {datetime.now().strftime('%H:%M:%S')}")
        print("=" * 40)
        
        counter_data = self.get_counter_data()
        
        for link, tunnels in self.stats_logger._link_tunnels().items():
            total = sum(counter_data.get(tid, 0) for tid, _ in tunnels)
            if total > 0:
                print(f"{link}: {total} packets")
    
    def start_monitoring(self):
        """Start the statistics monitoring loop"""
        try:
            while True:
                counter_data = self.get_counter_data()
                self.stats_logger.record_stats(counter_data)
                self.show_stats()
                sleep(2)
        except KeyboardInterrupt:
            print("\nStopping controller...")
        finally:
            self.stats_logger.close_logs()
            ShutdownAllSwitchConnections()


def main():
    parser = argparse.ArgumentParser(description='P4 Tunnel Controller')
    parser.add_argument('--p4info', help='Path to p4info proto in text format',
                        type=str, default='build/advanced_tunnel.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='Path to BMv2 JSON file',
                        type=str, default='build/advanced_tunnel.json')
    args = parser.parse_args()

    p4info_helper = p4runtime_lib.helper.P4InfoHelper(args.p4info)
    controller = P4Controller(p4info_helper, args.bmv2_json)
    
    controller.connect_switches()
    controller.setup_forwarding()
    controller.start_monitoring()


if __name__ == '__main__':
    main()